#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <errno.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netinet/tcp.h>
#include <pthread.h>
#include "test_utils.h"

static int set_tcp_no_delay(int fd)
{
    int flag = 1;
    return setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &flag, sizeof(flag));		
}

long long rtt;
long long count;
long long bytes;

void* print_stat(void*)
{
    long long sec = 0;
    while (true)
    {
        sleep(1);
        sec++;

        long long n_rtt = rtt;
        rtt = 0;

        long long n_count = count;
        count = 0;
        
        long long n_bytes = bytes;
        bytes = 0;

        printf("%lld: \t count: %lld \t latency(us): %lld \t thr(MB): %lld\n",
            sec,
            n_count,
            n_count == 0 ? 0 : n_rtt/(2*n_count),
            n_bytes/1024/1024);
    }

    return 0;
}

int main(int argc, char** argv)
{
    const int max_msg_size = 2*1024*1024;
    int msg_size = 4096;
    const char* host = "127.0.0.1";

    if (argc == 3)
    {
        host = argv[1];
        msg_size = atoi(argv[2]);
    }

    if (msg_size < 8)
        msg_size = 8;

    if (msg_size > max_msg_size)
        msg_size = max_msg_size;

    printf("host=%s, msg_size=%d\n", host, msg_size);

    int fd = socket(AF_INET, SOCK_STREAM, 0);
    errno_assert(fd != -1);

    sockaddr_in addr  = {};
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = inet_addr(host);
    addr.sin_port = htons(2650);

    printf("begin to connect to %s\n", host);
    int rc = connect(fd, (sockaddr*)&addr, sizeof(addr));
    errno_assert(rc == 0);
    printf("connect to %s successfully. start send...\n", host);

    rc = set_tcp_no_delay(fd);
    errno_assert(rc == 0);

    pthread_t tid;
    rc = pthread_create(&tid, 0, print_stat, 0);
    posix_assert(rc == 0);

    char buf[max_msg_size] = {};

    while (true)
    {
        long long us = now_us();

        int len = send(fd, buf, msg_size, 0);
        errno_assert(len == msg_size);

        len = recv(fd, buf, msg_size, MSG_WAITALL);
        errno_assert(len == msg_size || len == 0);
        if (len <= 0)
            break;

        rtt += (now_us() - us);
        count++;
        bytes += msg_size;
    }

    printf("send completed, close connection now.\n");

    rc = close(fd);
    errno_assert(rc == 0);	
	return 0;
}

